# -*- coding: utf-8 -*-
# @Time : 2021/11/1 16:01
# @Author : Ming
# @FileName: LSTM.py
# @Software: PyCharm

import torch


# parameters
num_class = 4
input_size = 4
hidden_size = 8
embedding_size = 10
num_layers = 2
batch_size = 1
seq_len = 5

# 通过RNN实现hello转换为ohlol
idx2char = ['e', 'h', 'l', 'o']
x_data = [[1, 0, 2, 2, 3]] # (batch, seq_len),这里变成了2个维度
y_data = [3, 1, 2, 3, 2] # (batch * seq_len)

inputs = torch.LongTensor(x_data)   # (batch, seq_len),这里变成了2个维度
labels = torch.LongTensor(y_data)   # (batch * seq_len)

# 设计模型  可以和RNN+embedding层的模型进行对比
class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.emb = torch.nn.Embedding(input_size, embedding_size)
        # input of RNN:(batch,seqlen,embeddingsize) output of RNN:(batch,seqlen,hiddensize)
        # LSTM更优秀的RNN,能够突破鞍点（减少梯度消失的问题），但是由于算法计算层数较多，时间复杂度较高
        self.lstm = torch.nn.LSTM(input_size=embedding_size,
                                  hidden_size=hidden_size,
                                  num_layers=num_layers,
                                  batch_first=True)
        # input of FC:(batch,seqlen,hiddensize) output of FC:(batch,seqlen,numclass)
        self.fc = torch.nn.Linear(hidden_size, num_class)

    def forward(self, x):
        hidden = torch.zeros(num_layers, x.size(0), hidden_size)
        c = torch.zeros(num_layers, x.size(0), hidden_size)
        x = self.emb(x) # (batch, seqlen, embeddingsize)
        x, _ = self.lstm(x, (hidden, c))
        x = self.fc(x)
        return x.view(-1, num_class)   # Reshape out to : (seqlen*batchsize, num_class)

model = Model()



net = Model()
criterion = torch.nn.CrossEntropyLoss() # 交叉熵
optimizer = torch.optim.Adam(net.parameters(), lr=0.05) # 优化器

# 训练循环开始
for epoch in range(15):
    optimizer.zero_grad()
    # RNN是直接获取所有输入数据
    outputs = net(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

    _, idx = outputs.max(dim=1)
    idx = idx.data.numpy()
    print('Predicted: ', ''.join([idx2char[x] for x in idx]), end='')
    print(', Epoch [%d/15] loss = %.3f' % (epoch + 1, loss.item()))